#/usr/bin/python
# -*- coding: utf-8 -*-

from edge import *

from numpy import argmin, argsort
import scipy.sparse as sp

class LocalBoundary(object):
    ''' 
    This class defines a boundary between two domains by linking one edge of a domain to one edge of another.
    It manages the nonconforming meshes by the use of the mortar method. Therefore a virtual boundary is discretized 
    and linked to each constitutive edges.
    --------------. <-> . <-> .--------------
      domain 1    .     
                  . <-> . <->      domain 2
                  . 
    --------------. <-> . <-> .--------------
               Edge1   B12  Edge2
    
    '''
    
    def __init__(self,domain1=None,edge1=None,domain2=None,edge2=None):
        '''
        Boundary constructor:
        @param domain1: 
        @param edge1:   
        @param domain2:
        @param edge2:
        
        .neighbours: list of code identifiers of associated domains
        .commonEdges: list of linked edges
        .binaryAssamblyOperators: list of binary assembly operators from 
        
        '''
        self.neighbours = [domain1,domain2]
        self.commonEdges = [edge1,edge2]
        #self.edgeTypes
        self.bmap=[None,None]
        self.binaryAssemblyOperators = []
        self.signedAssemblyOperators = []
        self.constraints = []
        self.mortarOperators = []
        self.traceOperators = []
        self.interpMesh = []

    def enrichBoundary(self):
        pass
    
    @staticmethod
    def buildAssemblyOperators(boundarymesh=None, edge1mesh=None, edge2mesh=None, mortars=[None,None]):
        '''
            Builds local assembly operators.
            @param boundarymesh: InterpMesh, optional.
            @param edge1mesh: InterpMesh, optional.
            @param edge2mesh: InterpMesh, optional.
            @param mortars: list of mortars operators, optional.
        '''
        binaryAOps = [None,None]
        signedAOps = [None,None]
        
#        if not boundarymesh:
#            boundarymesh = self.interpMesh
#            
#        if not edge1mesh:
#            try:
#                edge1mesh = self.commonEdges[0].interpMeshes[0]
#            except:
#                raise ValueError('Unable to get a mesh for the first edge.')
#            
#        if not edge2mesh:
#            try:
#                edge2mesh = self.commonEdges[1].interpMeshes[0]
#            except:
#                raise ValueError('Unable to get a mesh for the second edge.')
        
        if None in mortars and LocalBoundary._are_matching_meshes(edge1mesh, edge2mesh):
            bmap = LocalBoundary.buildBoundaryMap(boundarymesh, edge1mesh, edge2mesh)
            
            # number of degrees of freedom on boundary
            bnodes = len(bmap[0])
            NDOFS_PER_NODE = boundarymesh.ndofs_per_node
            bdofs = bnodes*NDOFS_PER_NODE
        
            # initialize local assembly operators
            binaryAOps = [sp.lil_matrix((bdofs,bdofs)),sp.lil_matrix((bdofs,bdofs))]
            signedAOps = [sp.lil_matrix((bdofs,bdofs)),sp.lil_matrix((bdofs,bdofs))]
        
            # local order of boundary nodes
            BMAP = [None,None]
            BMAP[0] = argsort(bmap[0])
            BMAP[1] = argsort(bmap[1])
            
            # fill assembly operators: local numbering of boundary virtual nodes is 0,1,2,... bnodes 
            # Hence, assembly operator maps local boundary node of domain to local boundary virtual node 
            
            for nd in xrange(bnodes):
                for d in xrange(NDOFS_PER_NODE):
                    binaryAOps[0][ nd*NDOFS_PER_NODE + d, BMAP[0][nd]*NDOFS_PER_NODE + d] = 1 
                    binaryAOps[1][ nd*NDOFS_PER_NODE + d, BMAP[1][nd]*NDOFS_PER_NODE + d] = 1
            
            signedAOps[0]=binaryAOps[0]
            signedAOps[1]=-binaryAOps[1]
            
        else:
            raise NotImplementedError
        
        return binaryAOps, signedAOps
    
    @staticmethod    
    def buildBoundaryMap(boundarymesh=None, edge1mesh=None, edge2mesh=None):
        '''
            Builds the mapping of nodes of edges towards nodes of boundary in case of matching meshes.
            @param boundarymesh: InterpMesh, optional.
            @param edge1mesh: InterpMesh, optional.
            @param edge2mesh: InterpMesh, optional.
        '''
        nnodes = len(boundarymesh.nodesCoo)
        bmap = [[0]*nnodes,[0]*nnodes]
        
        # edge1:
        for bnid,ndCoo in enumerate(boundarymesh.nodesMap['C2COO']):
             try:
                 nid1 = edge1mesh.nodesMap["COO2C"][ndCoo]
             except KeyError:
                 raise KeyError('Non matching node: '+str(bnid)+' at position '+ndCoo)
             try:
                 nid2 = edge2mesh.nodesMap["COO2C"][ndCoo]
             except:
                 raise KeyError('Non matching node: '+str(bnid)+' at position '+ndCoo)
             
             bmap[0][bnid]=nid1
             bmap[1][bnid]=nid2
                 
        return bmap
        
    def buildMortarOperators(self, boundarymesh=None, edge1mesh=None, edge2mesh=None):
        
        raise NotImplementedError('No mortar operators construction for now.')
    
    @staticmethod
    def setInterpMesh(edge1mesh=None, edge2mesh=None):
        '''
            Defines the mesh of the local boundary and builds the mortar operators if 
            the meshes of associated edges are non-matching.
            Inputs:
            @param edge1mesh: InterpMesh, optional. The mesh of the first edge.
            @param edge2mesh: InterpMesh, optional. The mesh of the second edge.
            Outputs:
            @return imesh: InterpMesh, mesh of the local boundary.
            @return mortars: list of mortar operators.
        '''
        # Initialize outputs:
        imesh = None
        mortars = [None,None]
        
#        if not edge1mesh:
#            try:
#                edge1mesh = self.commonEdges[0].interpMeshes[0]
#            except:
#                raise ValueError('Unable to get a mesh for the first edge.')
#            
#        if not edge2mesh:
#            try:
#                edge2mesh = self.commonEdges[1].interpMeshes[0]
#            except:
#                raise ValueError('Unable to get a mesh for the second edge.')
            
        # meshes of edges are matching:
        if LocalBoundary._are_matching_meshes(edge1mesh,edge2mesh):
            imesh = InterpMesh(**edge1mesh.__dict__)
        # meshes of edges are not matching:
        else:
            raise NotImplementedError('Non-matching edges have not been accounted for yet.')
            #TODO:
            # define an interpMesh
            i = argmin([len(edge1mesh.nodesCoo),len(edge2mesh.nodesCoo)])
            edgesmesh=[edge1mesh,edge2mesh]
            imesh = InterpMesh(**edgesmesh[i].__dict__)
            
            # calculate the mortar operators:
            mortars = LocalBoundary.buildMortarOperators(imesh,edge1mesh,edge2mesh)
        
        return imesh, mortars
    
    @staticmethod
    def _are_matching_meshes(mesh1, mesh2):
        
        matching = False
        
        nnodes = [len(mesh1.nodesCoo), len(mesh2.nodesCoo)]
        
        if nnodes[0]==nnodes[1]:
            matching = True
            for ndCoo in mesh1.nodesMap["COO2C"].keys():
                try:
                    mesh2.nodesMap["COO2C"][ndCoo]
                except KeyError:
                    matching = False
                    break
        
        return matching
        
    
class LocalBoundaryDecorator(LocalBoundary):
    '''
        Decorator of instances of class LocalBoundary(). 
    '''
    def __init__(self,localboundary): 
        kwargs = localboundary.__dict__
        super(LocalBoundary,self).__init__(kwargs['neighbours'][0],
                                           kwargs['commonEdges'][0],
                                           kwargs['neighbours'][1],
                                           kwargs['commonEdges'][1])
        
class PrimalBoundary(LocalBoundary):
    '''
        Primal boundary 
    '''

    def __init__(self):
        
        super(PrimalBoundary,self).__init__()
        
        
class DualBoundary(LocalBoundary):
    '''
        Dual boundary 
    '''

    def __init__(self):
        
        super(LocalBoundary,self).__init__()
        
class MixedBoundary(LocalBoundary):
    '''
        Mixed boundary 
    '''

    def __init__(self):
        
        super(LocalBoundary,self).__init__()

btype_to_modelClass = {
    'P': PrimalBoundary(),
    'D': DualBoundary(),
    'M': MixedBoundary(),
}